# model.py
import pymysql
from pymysql import Error

class DatabaseModel:
    def __init__(self):
        self.connection = None

    def connect(self, host, port, user, password, database=None):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(
                host=host,
                port=int(port),
                user=user,
                password=password,
                database=database,
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor
            )
            return True
        except Error as e:
            raise Exception(f"连接失败: {str(e)}")

    def test_connection(self, host, port, user, password):
        """测试数据库连接"""
        try:
            temp_conn = pymysql.connect(
                host=host,
                port=int(port),
                user=user,
                password=password,
                connect_timeout=5
            )
            temp_conn.close()
            return True
        except Error as e:
            raise Exception(f"连接测试失败: {str(e)}")

    def execute_query(self, query, params=None):
        """执行SQL查询"""
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(query, params or ())
                if query.strip().lower().startswith("select"):
                    return cursor.fetchall()
                self.connection.commit()
                return True
        except Error as e:
            self.connection.rollback()
            raise Exception(f"查询执行失败: {str(e)}")

    def get_tables(self):
        """获取数据库所有表"""
        return self.execute_query("SHOW TABLES")

    def close(self):
        """关闭连接"""
        if self.connection:
            self.connection.close()